import numpy as np
import matplotlib.pyplot as plt
from init_param import UAVParameters
from uav_model import UAVModel
from common_function import rotation_matrix, angle_body2world
from PIDController import posion_pos_controller
from LQRController import LQRController
from MPCController import MPCController
from square_signal import generate_ramp_signal
from manipulator_traj import man_desired,traj


class UAVSimulation:
    def __init__(self, params, dynamics, position_controller, pose_controller, controller_type):
        self.params = params
        self.dynamics = dynamics
        self.controller_type = controller_type
        if controller_type == "PID":
            self.x_controller = position_controller[0]
            self.y_controller = position_controller[1]
            self.height_controller = position_controller[2]
        else:
            self.position_controller = position_controller[0]
            self.velocity_controller = position_controller[1]
        self.roll_controller = pose_controller[0]
        self.pitch_controller = pose_controller[1]
        self.yaw_controller = pose_controller[2]
        self.position_history = []
        self.setpoint_history = []
        self.manipulator_history = []

    def desired_euler_angles(self, eax, eay, psi):
        """
        Convert the desired horizontal acceleration to the desired Euler angle
        Input:
          eax, eay: desired horizontal acceleration
        Output:
          phi, theta: desired roll and pitch
        """
        g = 9.8  # 重力加速度 (m/s^2)

        phi = (-np.sin(psi) * eax + np.cos(psi) * eay) / g
        theta = (-np.cos(psi) * eax - np.sin(psi) * eay) / g

        return phi, theta

    def motor_mixer(self, roll, pitch, yaw, thrust):
        """
        Control allocation. The quadrotor type is X-configuration,
        and the airframe is as follows:
        3↓   1↑
          \ /
          / \
        2↑   4↓

        Input:
          roll, pitch, yaw: attitude controller output.
          thrust: total thrust.

        Output:
          M1, M2, M3, M4: motor speed commands.
        """
        idle_PWM = 1000
        scale = 1000

        M1 = (thrust - roll + pitch + yaw) * scale + idle_PWM
        M2 = (thrust + roll - pitch + yaw) * scale + idle_PWM
        M3 = (thrust + roll + pitch - yaw) * scale + idle_PWM
        M4 = (thrust - roll - pitch - yaw) * scale + idle_PWM

        return M1, M2, M3, M4

    def pwm_to_rpm(self, pwm):
        """
        Convert PWM signal to RPM. This is a simple linear approximation.
        Input:
          pwm: PWM signal value

        Output:
          rpm: motor speed in RPM
        """
        min_pwm = 1000
        max_pwm = 2000

        # 限制PWM范围
        pwm = np.clip(pwm, min_pwm, max_pwm)

        # 线性插值计算RPM
        rpm = np.clip((pwm - min_pwm) / (max_pwm - min_pwm), -1, 1)
        return rpm

    def run(self, square_corners, duration, dt):
        DEG2RAD = self.params.DEG2RAD
        posE = self.params.ModelInit_PosE.astype(np.float64)
        velB = self.params.ModelInit_VelB.astype(np.float64)
        velE = self.params.ModelInit_VelB.astype(np.float64)
        angEuler = self.params.ModelInit_AngEuler
        rateB = self.params.ModelInit_RateB
        alpha_B = self.params.ModelInit_RateB
        rateE = self.params.ModelInit_RateB
        theta_m = self.params.ModelInit_theta_m
        theta_m_dot = self.params.ModelInit_theta_m
        theta_m_ddot = self.params.ModelInit_theta_m
        ModelParam_motorCr = self.params.ModelParam_motorCr
        ModelParam_motorWb = self.params.ModelParam_motorWb

        state = [velB, angEuler, rateB]
        man_state = [theta_m,theta_m_dot,theta_m_ddot]
        omega_state =[rateB,alpha_B]
        current_corner_index = 0
        time_steps = int(duration / dt)
        num_steps = time_steps // (len(square_corners) - 1)
        # next_corner_time = time_steps  // (len(square_corners)-1)
        input_signal = generate_ramp_signal(square_corners, num_steps)
        traj_theta_1 = traj(0, duration, 0, 20*np.pi/180)
        traj_theta_2 = traj(0, duration, 0, 30*np.pi/180)

        for step in range(time_steps):
            setpoint = input_signal[step]
            eax, eay, altitude_output = None, None, None
            print('setpoint:', setpoint)
            theta_m, theta_m_dot, theta_m_ddot = man_desired(traj_theta_1,traj_theta_2,step*dt)
            man_state = [theta_m, theta_m_dot, theta_m_ddot]
            if self.controller_type == 'PID':
                eax = self.x_controller.update(outer_setpoint=setpoint[0], outer_measurement=posE[0],
                                               inner_measurement=velE[0], dt=dt)
                eay = self.y_controller.update(outer_setpoint=setpoint[1], outer_measurement=posE[1],
                                               inner_measurement=velE[1], dt=dt)
                altitude_output = self.height_controller.update(outer_setpoint=setpoint[2], outer_measurement=posE[2],
                                                                inner_measurement=velE[2], dt=dt)
            else:
                vel_setpoint = self.position_controller.update(setpoint, posE)
                # 速度环LQR控制
                acc_setpoint = self.velocity_controller.update(vel_setpoint, velE)
                if self.controller_type == "MPC":
                    eax, eay, altitude_output = acc_setpoint[0][0], acc_setpoint[1][0], acc_setpoint[2][0]
                elif self.controller_type == "LQR":
                    eax, eay, altitude_output = acc_setpoint[0], acc_setpoint[1], acc_setpoint[2]
                else:
                    print("unknown controller")
            phi, theta = self.desired_euler_angles(eax, eay, angEuler[2])
            phi = np.clip(phi, -35 * DEG2RAD, 35 * DEG2RAD)
            theta = np.clip(theta, -35 * DEG2RAD, 35 * DEG2RAD)
            print('eax,eay:', eax, eay)
            print('phi,theta:', phi, theta)
            roll_output = np.clip(self.roll_controller.update(outer_setpoint=phi, outer_measurement=angEuler[0],
                                                              inner_measurement=rateE[0], dt=dt), -1, 1)
            pitch_output = np.clip(self.pitch_controller.update(outer_setpoint=theta, outer_measurement=angEuler[1],
                                                                inner_measurement=rateE[1], dt=dt), -1, 1)
            yaw_output = np.clip(self.yaw_controller.update(outer_setpoint=(step // num_steps) * 0.5 * 0.0174533,
                                                            outer_measurement=angEuler[2],
                                                            inner_measurement=rateE[2], dt=dt), -1, 1)
            thrust = np.clip(-altitude_output + self.params.THR_HOVER, 0.05, 0.9)
            print('roll_output,pitch_output,yaw_output,altitude_output:', roll_output, pitch_output, yaw_output,
                  altitude_output)
            M1, M2, M3, M4 = self.motor_mixer(roll_output, pitch_output, yaw_output, thrust)
            print('M1,M2,M3,M4:', M1, M2, M3, M4)

            w_target = (np.array(
                [self.pwm_to_rpm(M1), self.pwm_to_rpm(M2), self.pwm_to_rpm(M3),
                 self.pwm_to_rpm(M4)]) * ModelParam_motorCr +
                        np.array([ModelParam_motorWb, ModelParam_motorWb, ModelParam_motorWb, ModelParam_motorWb]))
            print('w_target:', w_target)
            state = self.dynamics.update_state(state,man_state,theta_m[0], w_target, dt)
            velB, angEuler, rateB = state
            # R = rotation_matrix(angEuler)
            R_angle = angle_body2world(angEuler)
            R = rotation_matrix(angEuler)
            rateE = R_angle @ rateB
            angEuler = angEuler + rateE * dt
            velE = R @ velB
            posE = posE + velE * dt

            state = velB, angEuler, rateB
            self.position_history.append(posE.copy())
            self.manipulator_history.append([theta_m[0],theta_m[1],theta_m_dot[0],theta_m_dot[1],theta_m_ddot[0],theta_m_ddot[1]])
            print('theta_m:',theta_m)
            print(
                f"Time: {step * dt:.2f}s, Position: {posE}, Velocity_world: {velE}, Angle: {angEuler}, Rate: {rateE},theta_m:{theta_m}")
            # if step % (next_corner_time-1) == 0 and step != 0:
            #     current_corner_index = (current_corner_index + 1) % len(square_corners)
            self.setpoint_history.append(setpoint.copy())

    def plot_trajectory(self):
        position_history = np.array(self.position_history)
        setpoint_history = np.array(self.setpoint_history)
        fig = plt.figure()
        ax = fig.add_subplot(121, projection='3d')
        ax.plot(position_history[:, 0], position_history[:, 1], -position_history[:, 2], label='UAV Trajectory')
        ax.set_xlabel('X Position (m)')
        ax.set_ylabel('Y Position (m)')
        ax.set_zlabel('Z Position (m)')
        ax.set_title('Trajectory of UAV ')
        ax.legend()

        ax = fig.add_subplot(122, projection='3d')
        ax.plot(setpoint_history[:, 0], setpoint_history[:, 1], -setpoint_history[:, 2], label='setpoint Trajectory')
        ax.set_xlabel('X Position (m)')
        ax.set_ylabel('Y Position (m)')
        ax.set_zlabel('Z Position (m)')
        ax.set_title(' Desired Trajectory ')
        ax.legend()


    def plot_state_error(self):
        if self.controller_type == 'PID':
            fig_position = plt.figure(figsize=(14, 4))
            self.x_controller.outer_controller.plot_error(fig_position, 0, title='x_position')
            self.y_controller.outer_controller.plot_error(fig_position, 1, title='y_position')
            self.height_controller.outer_controller.plot_error(fig_position, 2, title='height_position')
            fig_velocity = plt.figure(figsize=(14, 4))
            self.x_controller.inner_controller.plot_error(fig_velocity, 0, title='x_velocity')
            self.y_controller.inner_controller.plot_error(fig_velocity, 1, title='y_velocity')
            self.height_controller.inner_controller.plot_error(fig_velocity, 2, title='height_velocity')
        else:
            self.position_controller.plot_error('_position')
            self.velocity_controller.plot_error('_velocity')

        fig_angle = plt.figure(figsize=(14, 4))

        self.roll_controller.outer_controller.plot_error(fig_angle, 0, title='roll_angle')
        self.pitch_controller.outer_controller.plot_error(fig_angle, 1, title='pitch_angle')
        self.yaw_controller.outer_controller.plot_error(fig_angle, 2, title='yaw_angle')
        plt.tight_layout()
        fig_rate = plt.figure(figsize=(14, 4))

        self.roll_controller.inner_controller.plot_error(fig_rate, 0, title='roll_rate')
        self.pitch_controller.inner_controller.plot_error(fig_rate, 1, title='pitch_rate')
        self.yaw_controller.inner_controller.plot_error(fig_rate, 2, title='yaw_rate')
        plt.tight_layout()

        fig = plt.figure(figsize=(14, 4))
        ax = fig.add_subplot(131)
        ax.plot([i / 100 for i in range(len(self.manipulator_history))], [x[0] for x in self.manipulator_history], label='theta_m1')
        ax.plot([i / 100 for i in range(len(self.manipulator_history))], [x[1] for x in self.manipulator_history],
                label='theta_m2')
        ax.legend()
        ax = fig.add_subplot(132)
        ax.plot([i / 100 for i in range(len(self.manipulator_history))], [x[2] for x in self.manipulator_history],
                label='theta_m1_dot')
        ax.plot([i / 100 for i in range(len(self.manipulator_history))], [x[3] for x in self.manipulator_history],
                label='theta_m2_dot')
        ax.legend()
        ax = fig.add_subplot(133)
        ax.plot([i / 100 for i in range(len(self.manipulator_history))], [x[4] for x in self.manipulator_history],
                label='theta_m1_ddot')
        ax.plot([i / 100 for i in range(len(self.manipulator_history))], [x[5] for x in self.manipulator_history],
                label='theta_m2_ddot')
        ax.legend()
        plt.tight_layout()
        # plt.show()


if __name__ == '__main__':
    params = UAVParameters()
    # dynamics = UAVDynamics(params)
    muav = UAVModel(params)
    # controller = CascadePIDController(PIDController(3.0, 0.0, 0.15), PIDController(0.15, 0.0, 0.03))
    controller_type = "PID"
    # controller_type = "MPC"
    # controller_type = "LQR"
    position_controller = None
    if controller_type == 'PID':
        position_controller = posion_pos_controller(params).position_controller()
    elif controller_type == 'LQR':
        # 定义位置环的状态空间模型
        A_pos = np.zeros((3, 3))
        B_pos = np.eye(3)

        Q_pos = np.eye(3) * 16
        R_pos = np.eye(3) * 64

        # 定义速度环的状态空间模型
        A_vel = np.zeros((3, 3))
        B_vel = np.eye(3)

        Q_vel = np.eye(3) * 25
        R_vel = np.eye(3) * 49
        pos_controller = LQRController(A_pos, B_pos, Q_pos, R_pos)
        vel_controller = LQRController(A_vel, B_vel, Q_vel, R_vel)
        position_controller = [pos_controller, vel_controller]
    elif controller_type == 'MPC':
        # 定义离散位置环的状态空间模型
        A_pos = np.eye(3)
        B_pos = np.eye(3) * 0.01

        Q_pos = np.eye(3)*3.2
        R_pos = np.eye(3)

        # 定义离散速度环的状态空间模型
        # A_vel = np.zeros((3, 3))
        A_vel = np.eye(3)
        B_vel = np.eye(3) * 0.01

        Q_vel = np.eye(3) * 36
        R_vel = np.eye(3) * 16
        pos_controller = MPCController(A_pos, B_pos, None, 10, Q_pos, R_pos, -3, 3)
        vel_controller = MPCController(A_vel, B_vel, None, 10, Q_vel, R_vel, -1.5, 1.5)
        position_controller = [pos_controller, vel_controller]
    else:
        print('ERROR:Unknown Controller')

    pose_controller = posion_pos_controller(params).pos_controller()
    target_altitude = -5  # 目标高度
    square_corners = [
        [0, 0, 0],
        [0, 0, target_altitude],
        [5, 0, target_altitude],
        [5, 5, target_altitude],
        [0, 5, target_altitude],
        [0, 0, target_altitude],
        [0, 0, 0],
        [0, 0, 0]
    ]
    simulation = UAVSimulation(params, muav, position_controller=position_controller,
                               pose_controller=pose_controller, controller_type=controller_type)
    simulation.run(square_corners, 70, 0.01)
    simulation.plot_trajectory()
    simulation.plot_state_error()
    plt.show()
